param (
    [parameter(ParameterSetName="online", Mandatory=$true)]
    [switch] $FixupUsb,

    # vendor Ids of the various USB keys that i've worked with.
    [parameter(ParameterSetName="online")]
    $UsbVendorId = @("05dc", "090c", "0930", "1b1c"),

    [parameter(ParameterSetName="offline")]
	$DiskNumber,

    [parameter(ParameterSetName="offline", Mandatory=$true)]
    $ImagePath,

    [parameter(ParameterSetName="offline")]
    $Index = 1,

    [parameter(ParameterSetName="offline")]
    [switch]
    $WimBoot,

    [parameter(ParameterSetName="offline")]
    $EspLetter,

    [parameter(ParameterSetName="offline")]
    $OspLetter
	)

Set-StrictMode -Version 3
$ErrorActionPreference = "Stop"

function importKernel32 {
    # code copied and adapted from http://poshcode.org/4768
    # Author was Matthew Graeber (@mattifestation)
    # Under BSD 3-Clause license

    # Utilize P/Invoke in order to call QueryDosDevice. I prefer using 
    # reflection over Add-Type since it doesn't require compiling C# code.
    $DynAssembly = New-Object System.Reflection.AssemblyName('SysUtils')
    $AssemblyBuilder = [AppDomain]::CurrentDomain.DefineDynamicAssembly($DynAssembly, [Reflection.Emit.AssemblyBuilderAccess]::Run)
    $ModuleBuilder = $AssemblyBuilder.DefineDynamicModule('SysUtils', $False)

    # Define [Kenel32]
    $TypeBuilder = $ModuleBuilder.DefineType('Kernel32', 'Public, Class')

    # Define [Kernel32]::QueryDosDevice method
    $PInvokeMethod = $TypeBuilder.DefinePInvokeMethod('DefineDosDevice', 'kernel32.dll', ([Reflection.MethodAttributes]::Public -bor [Reflection.MethodAttributes]::Static), [Reflection.CallingConventions]::Standard, [UInt32], [Type[]]@([UInt32], [String], [String]), [Runtime.InteropServices.CallingConvention]::Winapi, [Runtime.InteropServices.CharSet]::Auto)
    $DllImportConstructor = [Runtime.InteropServices.DllImportAttribute].GetConstructor(@([String]))
    $SetLastError = [Runtime.InteropServices.DllImportAttribute].GetField('SetLastError')
    $SetLastErrorCustomAttribute = New-Object Reflection.Emit.CustomAttributeBuilder($DllImportConstructor, @('kernel32.dll'), [Reflection.FieldInfo[]]@($SetLastError), @($true))
    $PInvokeMethod.SetCustomAttribute($SetLastErrorCustomAttribute)

    # Define [Kernel32]::QueryDosDevice method
    $PInvokeMethod = $TypeBuilder.DefinePInvokeMethod('QueryDosDevice', 'kernel32.dll', ([Reflection.MethodAttributes]::Public -bor [Reflection.MethodAttributes]::Static), [Reflection.CallingConventions]::Standard, [UInt32], [Type[]]@([String], [Text.StringBuilder], [UInt32]), [Runtime.InteropServices.CallingConvention]::Winapi, [Runtime.InteropServices.CharSet]::Auto)
    $DllImportConstructor = [Runtime.InteropServices.DllImportAttribute].GetConstructor(@([String]))
    $SetLastError = [Runtime.InteropServices.DllImportAttribute].GetField('SetLastError')
    $SetLastErrorCustomAttribute = New-Object Reflection.Emit.CustomAttributeBuilder($DllImportConstructor, @('kernel32.dll'), [Reflection.FieldInfo[]]@($SetLastError), @($true))
    $PInvokeMethod.SetCustomAttribute($SetLastErrorCustomAttribute)

    return $TypeBuilder.CreateType()
}

function importCfgMgr32 {
    # code copied and adapted from http://poshcode.org/4768
    # Author was Matthew Graeber (@mattifestation)
    # Under BSD 3-Clause license

    # Utilize P/Invoke in order to call QueryDosDevice. I prefer using 
    # reflection over Add-Type since it doesn't require compiling C# code.
    $DynAssembly = New-Object System.Reflection.AssemblyName('SysUtils')
    $AssemblyBuilder = [AppDomain]::CurrentDomain.DefineDynamicAssembly($DynAssembly, [Reflection.Emit.AssemblyBuilderAccess]::Run)
    $ModuleBuilder = $AssemblyBuilder.DefineDynamicModule('SysUtils', $False)

    # Define [CfgMgr32]
    $TypeBuilder = $ModuleBuilder.DefineType('CfgMgr32', 'Public, Class')

    # Define [CfgMgr32]::CM_Locate_DevNode method
    $PInvokeMethod = $TypeBuilder.DefinePInvokeMethod('CM_Locate_DevNode', 'cfgmgr32.dll', ([Reflection.MethodAttributes]::Public -bor [Reflection.MethodAttributes]::Static), [Reflection.CallingConventions]::Standard, [UInt32], [Type[]]@([UInt32].MakeByRefType(), [String], [int]), [Runtime.InteropServices.CallingConvention]::Winapi, [Runtime.InteropServices.CharSet]::Auto)
    $DllImportConstructor = [Runtime.InteropServices.DllImportAttribute].GetConstructor(@([String]))

    # Define [CfgMgr32]::CM_Get_Parent method
    $PInvokeMethod = $TypeBuilder.DefinePInvokeMethod('CM_Get_Parent', 'cfgmgr32.dll', ([Reflection.MethodAttributes]::Public -bor [Reflection.MethodAttributes]::Static), [Reflection.CallingConventions]::Standard, [UInt32], [Type[]]@([UInt32].MakeByRefType(), [UInt32], [int]), [Runtime.InteropServices.CallingConvention]::Winapi, [Runtime.InteropServices.CharSet]::Auto)
    $DllImportConstructor = [Runtime.InteropServices.DllImportAttribute].GetConstructor(@([String]))

    # Define [CfgMgr32]::CM_Get_DeviceId method
    $PInvokeMethod = $TypeBuilder.DefinePInvokeMethod('CM_Get_Device_ID', 'cfgmgr32.dll', ([Reflection.MethodAttributes]::Public -bor [Reflection.MethodAttributes]::Static), [Reflection.CallingConventions]::Standard, [UInt32], [Type[]]@([UInt32], [Text.StringBuilder], [int], [int]), [Runtime.InteropServices.CallingConvention]::Winapi, [Runtime.InteropServices.CharSet]::Auto)
    $DllImportConstructor = [Runtime.InteropServices.DllImportAttribute].GetConstructor(@([String]))

    # Define [CfgMgr32]::CM_Disable_DevNode method
    $PInvokeMethod = $TypeBuilder.DefinePInvokeMethod('CM_Disable_DevNode', 'cfgmgr32.dll', ([Reflection.MethodAttributes]::Public -bor [Reflection.MethodAttributes]::Static), [Reflection.CallingConventions]::Standard, [UInt32], [Type[]]@([UInt32], [int]), [Runtime.InteropServices.CallingConvention]::Winapi, [Runtime.InteropServices.CharSet]::Auto)
    $DllImportConstructor = [Runtime.InteropServices.DllImportAttribute].GetConstructor(@([String]))

    # Define [CfgMgr32]::CM_Enable_DevNode method
    $PInvokeMethod = $TypeBuilder.DefinePInvokeMethod('CM_Enable_DevNode', 'cfgmgr32.dll', ([Reflection.MethodAttributes]::Public -bor [Reflection.MethodAttributes]::Static), [Reflection.CallingConventions]::Standard, [UInt32], [Type[]]@([UInt32], [int]), [Runtime.InteropServices.CallingConvention]::Winapi, [Runtime.InteropServices.CharSet]::Auto)
    $DllImportConstructor = [Runtime.InteropServices.DllImportAttribute].GetConstructor(@([String]))

    return $TypeBuilder.CreateType()
}

function defineDosDevice {
    [CmdletBinding(SupportsShouldProcess=$true)]
    param([parameter(Mandatory=$true)][string] $DeviceName,
          $TargetPath = $null,
          [switch] $Remove)

   $kernel32 = importKernel32

   [int] $flags = 8 # DDD_NO_BROADCAST_SYSTEM

    if ($Remove.IsPresent) {
        if ($TargetPath -eq $null) {
            $TargetPath = queryDosDevice -DeviceName $DeviceName
            $flags = $flags -bor 0x07 # DDD_RAW_TARGET_PATH | DDD_EXACT_MATCH_ON_REMOVE | DDD_REMOVE_DEFINITION
        } elseif ($TargetPath.StartsWith("\Device")) {
            $flags = $flags -bor 0x07 # DDD_RAW_TARGET_PATH | DDD_EXACT_MATCH_ON_REMOVE | DDD_REMOVE_DEFINITION
        } else {
            $flags = $flags -bor 0x06 # DDD_EXACT_MATCH_ON_REMOVE | DDD_REMOVE_DEFINITION
        }

        if ($PSCmdlet.ShouldProcess($DeviceName, "Removing definition of drive letter") -eq $false) {
            throw "Drive letter removal cancelled"
        }
    }

    if ($kernel32::DefineDosDevice($flags, $DeviceName, $TargetPath) -eq $false) {
        throw New-Object -TypeName System.ComponentModel.Win32Exception -ArgumentList @([System.Runtime.InteropServices.Marshal]::GetLastWin32Error())
    }
}

# function gle {return [System.Runtime.InteropServices.Marshal]::GetLastWin32Error()}

function queryDosDevice {
    param([parameter(Mandatory=$true)] $DeviceName)

    $kernel32 = importKernel32

    $Max = 65536
    $StringBuilder = New-Object System.Text.StringBuilder($Max)

    if ($kernel32::QueryDosDevice($DeviceName, $StringBuilder, $Max) -eq 0) {
        
        # if the drive letter isn't defined, return $null

        # Note - I was checking the error code to be sure the error was 
        #        expected, but that's turned out to be very fragile going
        #        through powershell.  
        #if (gle -eq 203) {
            return $null
        #}

        #throw New-Object -TypeName System.ComponentModel.Win32Exception -ArgumentList @([System.Runtime.InteropServices.Marshal]::GetLastWin32Error())
    }

    return $StringBuilder.ToString()
}

function queryFreeDriveLetters {
    param($Exclude)

    $freeLetters = @{}
    @('c'[0] .. 'z'[0]) | 
        ForEach-Object {"{0}:" -f [char]$_} |
        ForEach-Object {$freeLetters[$_] = $true}

    Get-WmiObject Win32_Volume | 
        Where-Object {$_.DriveLetter } |
        ForEach-Object {$freeLetters[$_.DriveLetter] = $false}

    $freeLetters.GetEnumerator() | ForEach-Object {if ($_.Value -eq $true) {$_.Key}} | Sort-Object
}

function updateHive {
    param(
        [parameter(Mandatory=$true, ParameterSetName="offline")]
        $Hive,
        [parameter(ParameterSetName="online")][switch]
        $Online,
        $VendorIDList
        )

    try {

        if ($PsCmdlet.ParameterSetName -eq "offline") {
            reg load HKLM\temp_hive $Hive | Write-Verbose
            $key = "HKLM:\temp_hive\ControlSet001\Control\UsbStor"
        } else {
            $key = "HKLM:\SYSTEM\ControlSet001\Control\UsbStor"
        }

        $VendorIDList | 
            ForEach-Object {
                $keyPath = Join-Path $key ($_ + "xxxx")
                Write-Verbose "Updating $keyPath to be non-removable"
                if (Test-Path -Path $keyPath) {Get-Item -Path $keyPath}
                else                          {New-Item -Path $keyPath}
            } |
            New-ItemProperty -Name DeviceHackFlags -Value 0x40 -ErrorAction SilentlyContinue | 
            Write-Verbose

    } finally {
        if ($PsCmdlet.ParameterSetName -eq "offline") {
            [gc]::collect()
            reg unload HKLM\temp_hive | Write-Verbose
        }
    }
}

function getDriveLetters {
    param ([ref][parameter(Mandatory=$true)] $EspLetter, 
           [ref][parameter(Mandatory=$true)] $OspLetter)

    [array] $freeLetters = queryFreeDriveLetters

    if ($EspLetter.Value -eq $null) {
        $EspLetter.Value = $freeLetters[0]
        $freeLetters = $freeLetters | Select-Object -Skip 1
    }

    if ($OspLetter.Value -eq $null) {
        $OspLetter.Value = $freeLetters[0]
    }

    # If the selected letters are already assigned, confirm 
    if (queryDosDevice -DeviceName $EspLetter.Value) {
        defineDosDevice -DeviceName $espLetter.Value -Remove -Confirm
    }

    if (queryDosDevice -DeviceName $OspLetter.Value) {
        defineDosDevice -DeviceName $OspLetter.Value -Remove -Confirm
    }

    Write-Verbose "Mounting ESP at $($espLetter.Value) and OS partition at $($ospLetter.Value)"
}

function getParentDevice($Device) {
    $cfgMgr32 = importCfgMgr32

    $dev = 0
    
    $result = $cfgMgr32::CM_Locate_DevNode([ref]$dev, $Device.DeviceID, 0)
    if ($result -ne 0) {throw "CfgMgr32 error ($result) calling CM_Locate_DevNode for $($Device.DeviceID)"}

    $parent = 0;
    $result = $cfgMgr32::CM_Get_Parent([ref]$parent, $dev, 0)
    if ($result -ne 0) {throw "CfgMgr32 error ($result) calling CM_Get_Parent for $($Device.DeviceID)"}

    $max = 1024
    $deviceId = New-Object System.Text.StringBuilder($Max)

    $result = $cfgMgr32::CM_Get_Device_ID($parent, $deviceId, $deviceId.Capacity, 0)
    if ($result -ne 0) {throw "CfgMgr32 error ($result) getting Device ID"}

    $pnpEntity = gwmi -Class Win32_PnpEntity | Where-Object {$_.DeviceId -eq $deviceId.ToString()}

    return $pnpEntity
}

function restartDevice($DeviceId) {
    $cfgMgr32 = importCfgMgr32

    $dev = 0
    
    $result = $cfgMgr32::CM_Locate_DevNode([ref]$dev, $DeviceID, 0)
    if ($result -ne 0) {throw "CfgMgr32 error ($result) calling CM_Locate_DevNode for $($Device.DeviceID)"}

    $result = $cfgMgr32::CM_Disable_DevNode($dev, 0)
    if ($result -ne 0) {throw "CfgMgr32 error ($result) disabling device $DeviceId"}

    $result = $cfgMgr32::CM_Enable_DevNode($dev, 0)
    if ($result -ne 0) {throw "CfgMgr32 error ($result) disabling device $DeviceId"}
}

function getUsbDiskParent($WmiDisk) {
    $pnpDisk = Get-WmiObject -Query "ASSOCIATORS OF {Win32_DiskDrive.DeviceId='$($wmiDisk.DeviceID)'} WHERE AssocClass=Win32_PnpDevice"
    return getParentDevice -Device $pnpDisk
}

function getDeviceHardwareId($PnpEntity) {
    if ($PnpEntity.HardwareID[0] -match "USB\\VID_(....)") {
        return $matches[1]
    } else {
        throw ("Unable to get VID from " + $pnpEntity.HardwareID[0])
    }
}

function selectTargetDisk {
    param([ref] $DiskNumber) 

    # get list of acceptable disks
    $disks = Get-Disk | 
        Sort-Object Number | 
        Where-Object {$_.OperationalStatus -eq "Online" -and $_.BusType -ne "RAID"}

    [int]$num = 0

    if ($DiskNumber.Value -eq $null) {

        # prompt the user to select a disk 

        do {

            $disks | 
                Select-Object Number,BusType,@{Name="Name";Expression={$_.FriendlyName}},@{Name="Size";Expression={"{0:n1} GB" -f ($_.Size / 1GB)}} |
                Format-Table -AutoSize |
                Out-Host

            $selected = Read-Host -Prompt "Please enter the number of the disk to install to:"
            if ($selected -eq 'q') {
                throw "disk selection aborted"
            } elseif ([int]::TryParse($selected, $DiskNumber) -eq $false) {
                Write-Host "Please enter a number"
                continue
            } elseif ($DiskNumber.Value -notin ($disks | Select-Object -ExpandProperty Number)) {
                Write-Host "Disk not found"
                continue
            } else {
                break
            }
        } while ($true) 

    } else {

        # examine the selected disk to make sure it's appropriate

        if ($DiskNumber.Value -notin ($disks | Select-Object -ExpandProperty Number)) {
            throw "Disk $DiskNumber is not an appropriate device"
        }
    }

    # That the disk number is used in the Name of the WMI object is an 
    # implementation detail.  But so far it's the best way i've found to 
    # get from the PS disk objects into the WMI objects, and the WMI
    # objects allow us to figure out if this is a Removable, USB disk.

    return gwmi -Query "SELECT * From Win32_DiskDrive WHERE Name='\\\\.\\PHYSICALDRIVE$($DiskNumber.Value)'"
}

function setupUsbDisk($WmiDisk)  {

    # apply the USBSTOR hack flag to the image
    # to make it look non-removable.  To do this we need the VID
    $parent = getUsbDiskParent -WmiDisk $wmiDisk
    $usbVid = getDeviceHardwareId -PnpEntity $parent

    # If the device is removable, then we can use the USBSTOR hack flags 
    # to make it non-removable.  Otherwise it's not an appropriate disk.

    if ($isRemovable) {

        Write-Host "Applying USB hack to selected disk & restarting it"
        updateHive -Online -VendorIDList @($usbVid)
        restartDevice -DeviceId $parent.DeviceId
        Start-Sleep -Seconds 5
    }

    return $usbVid
}


function applyImage {
    param (
	    $DiskNumber,

        [parameter(Mandatory=$true)]
        $ImagePath,

        $Index = 1,

        [switch]
        $WimBoot,

        $EspLetter,
        $OspLetter
	    )



    $scriptDir = $PSScriptRoot
    $espGuid =      "{C12A7328-F81F-11D2-BA4B-00A0C93EC93B}"
    $msrGuid =      "{E3C9E316-0B5C-4DB8-817D-F92DF00215AE}"
    $recoveryGuid = "{de94bba4-06d1-4d40-a16a-bfd50179d6ac}"

    # Get the WMI object for the target disk
    $wmiDisk = selectTargetDisk ([ref] $DiskNumber)

    # Does the device report that it can support removable media?
    $isRemovable = ($wmiDisk.Capabilities -contains 7)
    $isUsb = ($wmiDisk.InterfaceType -eq "USB")
    $usbVendorId = $null

    # If it's a removable USB disk then fixup the registry and restart the disk
    if ($isUsb) {
        $UsbVendorId = @(setupUsbDisk $wmiDisk)
    }

    $wim = Get-Item -Path $ImagePath

    if ($wim -eq $null) {
        throw "Wim $wim not found"
    }

    $efiBootPath="fs0:"

    $dismCmd = "dism.exe"

    $esp,$osp = $EspLetter,$OspLetter
    getDriveLetters -EspLetter ([ref] $esp) -OspLetter ([ref] $osp)

    # Note - these flags are useful when trying to debug the latter part
    #        of the script, so you can bypass the expensive image writing
    #        step.
    $reformatDisk = $true
    $applyImage = $true

    if ($reformatDisk) {

        Clear-Disk -Number $DiskNumber -RemoveData -RemoveOEM -ErrorAction Continue
        Initialize-Disk -Number $DiskNumber -PartitionStyle GPT

        # create the ESP
        New-Partition -DiskNumber $DiskNumber -GptType $espGuid -Size 100MB
        Add-PartitionAccessPath -DiskNumber $DiskNumber -PartitionNumber 2 -AccessPath $esp
        "Y" | format.com $esp /fs:FAT32 /Q

        # Create the MSR
        New-Partition -DiskNumber $DiskNumber -GptType $msrGuid -Size 256MB

        # Create the data partition
        $p = New-Partition -DiskNumber $DiskNumber -UseMaximumSize

        Format-Volume -Partition $p -FileSystem NTFS -Force -NewFileSystemLabel "System" -Confirm:$false
        Add-PartitionAccessPath -InputObject $p -AccessPath $osp
    } else {
        Add-PartitionAccessPath -DiskNumber $DiskNumber -PartitionNumber 2 -AccessPath $esp
        Add-PartitionAccessPath -DiskNumber $DiskNumber -PartitionNumber 4 -AccessPath $osp
    }

    # apply the Windows Image to the data partition

    if ($applyImage) {
        if ($WimBoot.IsPresent -eq $false) {
            write-host "Applying Windows Image"
            & $dismCmd /Apply-Image /ImageFile:$ImagePath /Index:$Index /ApplyDir:$osp\
        } else {
            New-Item -Path "$osp\WIMs" -ItemType Directory > $null
    
            write-host "Copying Image file"
            $copy = Copy-Item -Path $ImagePath -Destination "$osp\WIMs" -PassThru
    
            write-host "Applying Windows Image"
            & $dismCmd /Apply-Image /ImageFile:$($copy.FullName) /Index:$Index /ApplyDir:$osp\ /WIMBoot
        }
    }

    # setup the disk for booting & testing
    if ($reformatDisk) {
        bcdboot.exe $osp\Windows /s $esp /f UEFI /v
    }
    bcdedit.exe -store "$esp\efi\microsoft\boot\bcd" -set "{default}" testsigning on

    # Update the hive on the target to make the USB device we're provisioning non-removable
    # so that we can boot from it.
    updateHive -Hive $osp\Windows\system32\CONFIG\SYSTEM -VendorIdList $UsbVendorId

    # Add the startup script to the root
    $efiBootPath + (Split-Path -Resolve -NoQualifier $esp\EFI\Boot\boot*) > $esp\startup.nsh
    
    # clean up the drive letters we created
    if ($EspLetter -eq $null) {
        defineDosDevice -DeviceName $esp -Remove -Confirm:$false
    }

    if ($OspLetter -eq $null) {
        defineDosDevice -DeviceName $osp -Remove -Confirm:$false
    }
}

switch ($PSCmdlet.ParameterSetName) {
    "online" {
        updateHive -Online -VendorIDList $UsbVendorId
        break;
    }
    "offline" {
        applyImage @PSBoundParameters
        break;
    }
}
